
def partition1(vec, predicate):
    N = len(vec)
    PARENT=0
    RANK=1

    nodes = [] #(N*2);

    # The first O(N) pass: create N single-vertex trees
    for i in range(N):
        nodes.append([0, 0])
        nodes[i][PARENT]=-1
        nodes[i][RANK] = 0

    # The main O(N^2) pass: merge connected components
    for i in range(N):
        root = i

        # find root
        while( nodes[root][PARENT] >= 0 ):
            root = nodes[root][PARENT]

        for j in range(N):
            if( i == j or not predicate(vec[i], vec[j])):
                continue

            root2 = j

            while( nodes[root2][PARENT] >= 0 ):
                root2 = nodes[root2][PARENT]

            if( root2 != root ):
                # unite both trees
                rank = nodes[root][RANK]
                rank2 = nodes[root2][RANK]
                if( rank > rank2 ):
                    nodes[root2][PARENT] = root
                else:
                    nodes[root][PARENT] = root2
                    nodes[root2][RANK] += rank == rank2
                    root = root2
                
                assert( nodes[root][PARENT] < 0 )

                k = j
                parent = 0

                # compress the path from node2 to root
                while( (nodes[k][PARENT]) >= 0 ):
                    parent = nodes[k][PARENT]
                    nodes[k][PARENT] = root
                    k = parent

                # compress the path from node to root
                k = i
                while( (nodes[k][PARENT]) >= 0 ):
                    parent = nodes[k][PARENT]
                    nodes[k][PARENT] = root
                    k = parent

    # Final O(N) pass: enumerate classes
    labels = []
    nclasses = 0

    for i in range(N):
        root = i
        while( nodes[root][PARENT] >= 0 ):
            root = nodes[root][PARENT]
        # re-use the rank as the class label
        if( nodes[root][RANK] >= 0 ):
            nclasses+=1
            nodes[root][RANK] = ~nclasses
        labels.append(~nodes[root][RANK])

    return nclasses, labels


def partition_split(vec, nclasses, labels, mincnt = 0):
    out = []
    for i in range(nclasses):
        out.append([])

    for i in range(len(vec)):
        j = labels[i]-1
        out[j].append(vec[i])

    myList1 = sorted(out, key=lambda i: len(i), reverse=True)
    if mincnt>0:
        myList2 = []
        for x in myList1:
            if len(x)>mincnt:
                myList2.append(x)
        
        return myList2

    return myList1


def partition2(vec, predicate, mincnt = 0):
    nclasses, labels = partition1(vec, predicate)
    myList1 = partition_split(vec, nclasses, labels, mincnt)
    return myList1
